import os
from PIL import Image
import torch
from torch.utils.data import Dataset
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import albumentations as A
from albumentations.pytorch import ToTensorV2
class Mdataset(Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.raw_dir = os.path.join(root_dir, 'raw')
        self.mask_dir = os.path.join(root_dir, 'label')
        self.transform = transform
    # Ensure the lists of raw and mask images are initialized
        self.raw_images = sorted(os.listdir(self.raw_dir))
        self.mask_images = sorted(os.listdir(self.mask_dir))

    def __len__(self):
        return len(self.raw_images)


    def __getitem__(self, idx):
        raw_path = os.path.join(self.raw_dir, self.raw_images[idx])
        mask_path = os.path.join(self.mask_dir, self.mask_images[idx])

        # Load raw image (color) using PIL
        raw_img = Image.open(raw_path).convert('RGB')  # Ensure it's in RGB format

        # Load mask (black and white) using PIL
        mask_img = Image.open(mask_path).convert('L')  # Ensure it's in grayscale (L mode)

        if self.transform:
            # Convert images to numpy arrays before passing to Albumentations
            raw_img = np.array(raw_img)
            mask_img = np.array(mask_img)

            augmented = self.transform(image=raw_img, mask=mask_img)
            raw_img = augmented['image']
            mask_img = augmented['mask']

        return raw_img, mask_img
transform = A.Compose([
    A.Resize(256, 256),
    A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    ToTensorV2()
])

def count_files_in_directory(directory):  #递归看数量
    file_count = 0
    for root, dirs, files in os.walk(directory):
        # Count the number of files in the current directory
        file_count = len(files)
        print(f"Directory: {root}, Files: {file_count}")
        for dir in dirs:
            # Recursive call to traverse subdirectories
            dir_path = os.path.join(root, dir)
            count_files_in_directory(dir_path)
def show_images(imgs, num_rows, num_cols, titles=None, scale=1.5):
    """Plot a list of images.

    Defined in :numref:`sec_fashion_mnist`"""
    figsize = (num_cols * scale, num_rows * scale)
    _, axes = plt.subplots(num_rows, num_cols, figsize=figsize)
    axes = axes.flatten()
    for i, (ax, img) in enumerate(zip(axes, imgs)):
        if torch.is_tensor(img):
            # Tensor Image
            ax.imshow(img.numpy())
        else:
            # PIL Image
            ax.imshow(img)
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        if titles:
            ax.set_title(titles[i])
    return axes

if __name__ == '__main__':
    if 1:
        # 定义图像的预处理
        transform = transforms.Compose([
            transforms.Resize((256, 256)),
            transforms.ToTensor(),
        ])
        root_dir='/home/jinyue/work/pcy241003/data/edeme/'

        dataset = Mdataset(root_dir, transform=transform)
        # 数据加载器
        dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=True)
        # 遍历数据集
        for images, masks in dataloader:
            print(f"Image batch shape: {images.shape}")
            print(f"Mask batch shape: {masks.shape}")
    if 0:

        # Replace with the directory you want to analyze
        directory_path = "pcy241003"
        count_files_in_directory(directory_path)
